#include "videorecorder_ffmpeg.h"
#include <QDebug>

VideoRecorder_ffmpeg::VideoRecorder_ffmpeg(QObject *parent)
    : BaseRecorder{parent}
{

}

bool VideoRecorder_ffmpeg::init()
{
    //
    if(!init_encode()){
        return false;
    }

    if(!init_file()){
        return false;
    }

    // 初始化packet
    m_pPacket = av_packet_alloc();

    // 初始化frame
    m_pFrame = av_frame_alloc();
    if (!m_pFrame) {
        fprintf(stderr, "Could not allocate video frame\n");
        exit(1);
    }

    m_pFrame->format = m_pCodecCtx->pix_fmt;
    m_pFrame->width  = m_pCodecCtx->width;
    m_pFrame->height = m_pCodecCtx->height;

    int ret = av_frame_get_buffer(m_pFrame, 32);

    if(ret < 0) {
        qDebug() << "init frame failed";
        return false;
    }

    m_bInited = true;
    return true;
}

bool VideoRecorder_ffmpeg::openFile(const QString &basename)
{
    m_sVideoBaseName = basename;

//    // 第一帧数据
    m_bFirstFrame = true;

    return true;
}

void VideoRecorder_ffmpeg::onImage(std::shared_ptr<QImage> img)
{
    if(!m_bInited) return;

    // 格式转换
    if(!image2Yuv420(img)){
        return;
    }

    // 编码
    if(!encode()){
        qDebug() << "cncode failed\n";
        return;
    }

    // 写文件
//    if(m_bFirstFrame){
//        if(writeHeader()){
//            qDebug() << "write head success!";
//        }
//       m_bFirstFrame = false;
//    }

//    writeFile();

}

bool VideoRecorder_ffmpeg::init_encode()
{
    m_pCodec = avcodec_find_encoder(m_codecId);
    if(nullptr == m_pCodec) {
        return false;
    }

    m_pCodecCtx = avcodec_alloc_context3(m_pCodec);
    if(nullptr == m_pCodecCtx) {
        return false;
    }

    m_pCodecCtx->bit_rate = 4000000;
    m_pCodecCtx->width = 1920;
    m_pCodecCtx->height = 1080;

    /* frames per second */
    m_pCodecCtx->time_base = (AVRational){1, 20};
    m_pCodecCtx->framerate = (AVRational){20, 1};
    m_pCodecCtx->gop_size = 10;
    m_pCodecCtx->max_b_frames = 1;
    m_pCodecCtx->pix_fmt = AV_PIX_FMT_YUV420P;
    if (m_pCodec->id == AV_CODEC_ID_H264)
            av_opt_set(m_pCodecCtx->priv_data, "preset", "slow", 0);

    int ret = avcodec_open2(m_pCodecCtx, m_pCodec, NULL);
    if (ret < 0) {
        qDebug() << "Could not open codec";
        return false;
    }

    return true;
}

bool VideoRecorder_ffmpeg::init_file()
{
    int ret = avformat_alloc_output_context2(&m_pFormatCtx, nullptr, nullptr, m_sVideoBaseName.toLatin1().data());
    if(ret < 0){
        return false;
    }
    m_pStream = avformat_new_stream(m_pFormatCtx, m_pCodec);

//    ret= avcodec_parameters_from_context(m_pStream->codecpar, m_pCodecCtx);
//    if(ret < 0){
//        qDebug() << "init parameters failed";
//        return false;
//    }

    m_pStream->time_base = m_pCodecCtx->time_base;
    m_pStream->id = m_pFormatCtx->nb_streams - 1;
    m_pStream->index = m_pFormatCtx->nb_streams - 1;
    m_pStream->codecpar->codec_tag = 0;
    m_pStream->codecpar->codec_type = m_pCodec->type;
    m_pStream->codecpar->codec_id = m_pCodec->id;
    m_pStream->codecpar->format =m_pCodecCtx->pix_fmt;
    m_pStream->codecpar->width = m_pCodecCtx->width;
    m_pStream->codecpar->height = m_pCodecCtx->height;
    m_pStream->codecpar->bit_rate = m_pCodecCtx->bit_rate;

    if (m_pFormatCtx->oformat->flags & AVFMT_GLOBALHEADER)
    {
        m_pFormatCtx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
    }

     qDebug() << "stream index:" << m_pStream->index;
    if(!(m_pFormatCtx->flags & AVFMT_NOFILE)){
        ret = avio_open(&m_pFormatCtx->pb, m_sVideoBaseName.toLatin1().data(), AVIO_FLAG_READ_WRITE);
        if(ret < 0){
            return false;
        }
    }
    return  writeHeader();
}

bool VideoRecorder_ffmpeg::writeHeader()
{
    int ret = avformat_write_header(m_pFormatCtx, nullptr);

    if(ret < 0){
        return false;
    }

    return true;
}

bool VideoRecorder_ffmpeg::encode()
{
    int ret = avcodec_send_frame(m_pCodecCtx, m_pFrame);
    if (ret < 0) {
        fprintf(stderr, "Error sending a frame for encoding\n");
        exit(1);
      }

    while (ret >= 0) {
        ret = avcodec_receive_packet(m_pCodecCtx, m_pPacket);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
            break;
        } else if (ret < 0) {
            return false;
//            break;
        }
        writeFile();
        av_packet_unref(m_pPacket);
    }




    return true;
}

bool VideoRecorder_ffmpeg::writeFile()
{
    m_pPacket->stream_index = 0;
//    qDebug() << "pts:" << m_pPacket->pts << ", dts:" << m_pPacket->dts << ",duration" << m_pPacket->duration;
    m_pPacket->pts = av_rescale_q_rnd(m_pPacket->pts, m_pCodecCtx->time_base, m_pStream->time_base, AV_ROUND_NEAR_INF);
//    m_pPacket->dts = av_rescale_q_rnd(m_pPacket->dts, m_pFormatCtx->streams[0]->time_base, m_pStream->time_base, AV_ROUND_NEAR_INF);
//    m_pPacket->duration = av_rescale_q_rnd(m_pPacket->duration, m_pFormatCtx->streams[0]->time_base, m_pStream->time_base, AV_ROUND_NEAR_INF);
//    m_pPacket->pos = -1;

//    qDebug() << "pts:" << m_pPacket->pts << ", dts:" << m_pPacket->dts << ",duration" << m_pPacket->duration;
    int ret = av_interleaved_write_frame(m_pFormatCtx, m_pPacket);
    if(ret < 0) {
        qDebug() << "wirte file failed";
        return false;
    }
    return true;
}

// 这里的img 只能接收rgb888格式
bool VideoRecorder_ffmpeg::image2Yuv420(std::shared_ptr<QImage> img)
{
    if(nullptr == m_pConvert){

    }

    const int in_linesize[1] = { img->bytesPerLine() };
    const uint8_t *in_data[1] = {img->bits()};
    m_pConvert = sws_getCachedContext(m_pConvert, img->width(),  img->height(), AV_PIX_FMT_RGB24, m_pCodecCtx->width, m_pCodecCtx->height, AV_PIX_FMT_YUV420P, 0,0,0,0);

    int ret = sws_scale(m_pConvert, in_data, in_linesize, 0, m_pFrame->height, m_pFrame->data, m_pFrame->linesize);
    if(ret < 0) {
        qDebug() << "rgb 2 yuv failed";
        return false;
    }

    m_pFrame->pts = m_iPts++;

    return true;
}
